#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
比较不同模型的准确性
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from predict_category import load_model, predict_category_with_subcategory

# 测试数据和预期结果
test_data = [
    ("星巴克咖啡", ("餐饮", "咖啡")),
    ("地铁出行", ("交通", "公共交通")),
    ("超市购物", ("购物", "日用品")),
    ("电影票", ("娱乐", "电影")),
    ("医院挂号", ("医疗", "门诊")),
    ("在线课程", ("教育", "在线教育")),
    ("房租支付", ("生活缴费", "房租")),
    ("水电费", ("生活缴费", "水电费")),
    ("外卖订单", ("餐饮", "外卖")),
    ("滴滴出行", ("交通", "打车")),
    ("京东购物", ("购物", "电子产品")),
    ("其他消费", ("其他", "其他")),
    ("肯德基快餐", ("餐饮", "快餐")),
    ("公交卡充值", ("交通", "公共交通")),
    ("淘宝购物", ("购物", "服饰")),
    ("游戏充值", ("娱乐", "游戏")),
    ("体检费用", ("医疗", "体检")),
    ("培训课程", ("教育", "培训")),
    ("物业费", ("生活缴费", "物业费")),
    ("燃气费", ("生活缴费", "燃气费"))
]

# 模型路径
models = {
    "原始模型": "./model_cache/bert_transaction_classifier",
    "改进模型": "./model_cache/improved_bert_transaction_classifier",
    # "剪枝模型": "./model_cache/bert_transaction_classifier_pruned",
    # "量化模型": "./model_cache/bert_transaction_classifier_quantized"
}

if __name__ == "__main__":
    print("模型准确性比较结果:")
    
    for model_name, model_path in models.items():
        try:
            print(f"\n测试 {model_name} ({model_path}):")
            model, tokenizer, model_type = load_model(model_path)
            
            correct = 0
            total = len(test_data)
            
            for text, expected in test_data:
                try:
                    primary, secondary, confidence = predict_category_with_subcategory(text, model, tokenizer, model_type)
                    actual = (primary, secondary)
                    
                    if actual == expected:
                        correct += 1
                        status = "✓"
                    else:
                        status = "✗"
                    
                    print(f"{status} '{text}' -> 预期: {expected}, 实际: {actual}")
                except Exception as e:
                    print(f"✗ '{text}' -> 错误: {e}")
            
            accuracy = correct / total * 100
            print(f"\n{model_name} 准确率: {correct}/{total} ({accuracy:.2f}%)")
        except Exception as e:
            print(f"加载 {model_name} 失败: {e}")